import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import KFold
from sklearn.metrics import accuracy_score, f1_score
#如果需要在美国建立一个娱乐设施工厂和一个服装厂为这些仓库供货，这两个厂应该分别建在哪里

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
pd.options.display.float_format = '{:.2f}'.format

#基于第二小问的存储仓库的城市编号
pos=[2,4,8,10,11,15,16,18,19,22,24,29,31,38,41,47,49,50]

#读取文件
df_1 = pd.read_excel('练习题一\B题-附件1.xls')
df_2 = pd.read_excel('练习题一\B题-附件2.xls')

#对洲的二个城市，直接取平均
df_2.loc[3,"服装消费"]=df_2.loc[4,"服装消费"]=df_2.loc[3,"服装消费"]/2
df_2.loc[3,"娱乐设施消费"]=df_2.loc[4,"娱乐设施消费"]=df_2.loc[3,"娱乐设施消费"]/2
df_2.loc[29,"服装消费"]=df_2.loc[30,"服装消费"]=df_2.loc[29,"服装消费"]/2
df_2.loc[29,"娱乐设施消费"]=df_2.loc[30,"娱乐设施消费"]=df_2.loc[29,"娱乐设施消费"]/2

# 补齐对称矩阵
for i in range(1, 51):
    for j in range(1, 51):
        if pd.isna(df_1[i][j]):
            df_1[i][j] = df_1[j][i]

#df_2列名：  城市编码    城市   州 州名简称     服装消费   娱乐设施消费   人口  消费税率  服装消费税率
# 方法一:分别寻找娱乐设施和服装厂的最优位置
def find_best_location(factory_type):
    best_city = None
    best_metric = -4858979 
    for city in range(1, 51):
        total_distance = 0
        total_consumption = 0
        for warehouse in pos:
            total_distance += df_1[city][warehouse]
            if factory_type == '娱乐':
                # 娱乐设施消费增加量
                total_consumption += df_2.loc[warehouse-1, "娱乐设施消费"] - df_2.loc[city-1, "娱乐设施消费"]
            else:
                # 服装消费增加量
                total_consumption += df_2.loc[warehouse-1, "服装消费"] - df_2.loc[city-1, "服装消费"]
        # 对于娱乐设施，我们希望消费增加量最大且距离较小
        # 对于服装厂，我们希望消费增加量最大且距离较小
        # 使用加权评分
        metric = total_consumption * 0.7 - total_distance * 0.3
        
        if factory_type == '娱乐':
            if metric > best_metric:
                best_metric = metric
                best_city = city
        else:
            if metric > best_metric:
                best_metric = metric
                best_city = city
        print(f"{best_city} {total_distance}  {total_consumption}")
    return best_city

# 分别寻找最优位置
best_entertainment = find_best_location('娱乐')
print("##############################")
best_clothing = find_best_location('服装')
print(f"最优选址：娱乐设施工厂在城市 {best_entertainment}，服装厂在城市 {best_clothing}")


#方法二：小数量级的应该改为随机森林
# 计算每个候选城市到所有仓库的总距离
def uty(i):
    total=0
    for a in pos:
        total+=df_1[i][a]
    return total

df_2['距离'] = 0
for city in range(1,51):
    df_2.loc[city-1 , '距离'] =uty(city)

# 定义特征和目标
features = df_2[[ "人口", "消费税率", "服装消费税率", "距离"]]
target_ent = df_2["娱乐设施消费"]  # 娱乐设施消费作为娱乐工厂的目标
target_clo = df_2["服装消费"]      # 服装消费作为服装厂的目标
# 处理可能的字符串数值(有NAN值的)
for col in features.columns:
    if features[col].dtype == object:
        features[col] = features[col].str.replace(',', '').astype(float)
        
# 数据标准化
scaler = MinMaxScaler()
scaled_features = scaler.fit_transform(features)

# 划分训练集和测试集
X_train, X_test, y_ent_train, y_ent_test = train_test_split(
    scaled_features, target_ent, test_size=0.2, random_state=42)
_, _, y_clo_train, y_clo_test = train_test_split(
    scaled_features, target_clo, test_size=0.2, random_state=42)

# 训练娱乐设施消费预测模型
rf_ent = RandomForestRegressor(n_estimators=150, random_state=42)
rf_ent.fit(X_train, y_ent_train)
ent_pred = rf_ent.predict(X_test)
print(f"娱乐设施消费模型测试集RMSE: {np.sqrt(mean_squared_error(y_ent_test, ent_pred)):.2f}")

# 训练服装消费预测模型
rf_clo = RandomForestRegressor(n_estimators=150, random_state=42)
rf_clo.fit(X_train, y_clo_train)
clo_pred = rf_clo.predict(X_test)
print(f"服装消费模型测试集RMSE: {np.sqrt(mean_squared_error(y_clo_test, clo_pred)):.2f}")

# 预测所有城市的消费潜力
all_ent_pred = rf_ent.predict(scaled_features)
all_clo_pred = rf_clo.predict(scaled_features)

# 结合距离因素计算综合得分
# 距离权重 - 可以根据实际情况调整
distance_weight = 0.3
consumption_weight = 0.7

# 标准化距离和消费预测值
scaled_distance = MinMaxScaler().fit_transform(df_2[['距离']])
scaled_ent_pred = MinMaxScaler().fit_transform(all_ent_pred.reshape(-1, 1))
scaled_clo_pred = MinMaxScaler().fit_transform(all_clo_pred.reshape(-1, 1))

# 计算综合得分
df_2['娱乐设施得分'] = consumption_weight * scaled_ent_pred - distance_weight * scaled_distance
df_2['服装厂得分'] = consumption_weight * scaled_clo_pred - distance_weight * scaled_distance

# 找出得分最高的城市
best_ent_city = df_2['城市编码'].iloc[df_2['娱乐设施得分'].argmax()]
best_clo_city = df_2['城市编码'].iloc[df_2['服装厂得分'].argmax()]

print("\n基于随机森林模型的最优选址结果:")
print(f"娱乐设施工厂最佳位置: 城市 {best_ent_city} - {df_2.loc[best_ent_city-1, '城市']}, {df_2.loc[best_ent_city-1, '州']}")
print(f"服装厂最佳位置: 城市 {best_clo_city} - {df_2.loc[best_clo_city-1, '城市']}, {df_2.loc[best_clo_city-1, '州']}")



# ============== 图像显示部分 ==============
plt.figure(figsize=(15, 8))

# 1. 仓库和工厂位置地图
plt.subplot(2, 2, 1)
plt.scatter(range(1,51), [1]*50, label='所有城市', color='gray', alpha=0.5)
plt.scatter(pos, [1]*len(pos), label='仓库位置', color='blue')
plt.scatter([best_entertainment, best_clothing], [1,1], 
            color=['red', 'green'], marker='*', s=200, 
            label=['娱乐工厂(方法一)', '服装工厂(方法一)'])
plt.scatter([best_ent_city, best_clo_city], [1.05,1.05], 
            color=['red', 'green'], marker='^', s=100, 
            label=['娱乐工厂(方法二)', '服装工厂(方法二)'])
plt.title('仓库和工厂位置分布')
plt.xlabel('城市编号')
plt.yticks([])
plt.legend()

# 2. 消费与距离关系图
plt.subplot(2, 2, 2)
plt.scatter(df_2['距离'], df_2['娱乐设施消费'], alpha=0.5, label='娱乐设施消费')
plt.scatter(df_2['距离'], df_2['服装消费'], alpha=0.5, label='服装消费')
plt.axvline(x=df_2.loc[best_entertainment-1, '距离'], color='red', linestyle='--', label='娱乐工厂(方法一)')
plt.axvline(x=df_2.loc[best_clothing-1, '距离'], color='green', linestyle='--', label='服装工厂(方法一)')
plt.axvline(x=df_2.loc[best_ent_city-1, '距离'], color='red', linestyle=':', label='娱乐工厂(方法二)')
plt.axvline(x=df_2.loc[best_clo_city-1, '距离'], color='green', linestyle=':', label='服装工厂(方法二)')
plt.title('消费水平与总运输距离关系')
plt.xlabel('到所有仓库的总距离')
plt.ylabel('消费水平')
plt.legend()

# 3. 随机森林特征重要性
plt.subplot(2, 2, 3)
features = ['人口', '消费税率', '服装消费税率', '距离']
ent_importance = rf_ent.feature_importances_
clo_importance = rf_clo.feature_importances_
x = range(len(features))
plt.bar(x, ent_importance, width=0.4, label='娱乐设施模型')
plt.bar([i+0.4 for i in x], clo_importance, width=0.4, label='服装模型')
plt.xticks([i+0.2 for i in x], features)
plt.title('随机森林特征重要性')
plt.legend()

# 4. 模型预测效果
plt.subplot(2, 2, 4)
plt.scatter(y_ent_test, ent_pred, alpha=0.5, label='娱乐设施')
plt.scatter(y_clo_test, clo_pred, alpha=0.5, label='服装')
plt.plot([min(y_ent_test.min(), y_clo_test.min()), max(y_ent_test.max(), y_clo_test.max())], 
         [min(y_ent_test.min(), y_clo_test.min()), max(y_ent_test.max(), y_clo_test.max())], 
         'k--')
plt.title('模型预测效果')
plt.xlabel('实际值')
plt.ylabel('预测值')
plt.legend()

plt.tight_layout()
plt.show()

# 输出最终结果对比
print("\n两种方法选址结果对比:")
print("方法一(暴力法):")
print(f"  娱乐工厂: 城市 {best_entertainment} - {df_2.loc[best_entertainment-1, '城市']}, {df_2.loc[best_entertainment-1, '州']}")
print(f"  服装工厂: 城市 {best_clothing} - {df_2.loc[best_clothing-1, '城市']}, {df_2.loc[best_clothing-1, '州']}")

print("\n方法二(随机森林模型):")
print(f"  娱乐工厂: 城市 {best_ent_city} - {df_2.loc[best_ent_city-1, '城市']}, {df_2.loc[best_ent_city-1, '州']}")
print(f"  服装工厂: 城市 {best_clo_city} - {df_2.loc[best_clo_city-1, '城市']}, {df_2.loc[best_clo_city-1, '州']}")

# 结果分析
if best_entertainment == best_ent_city and best_clothing == best_clo_city:
    print("\n两种方法选址结果一致，建议采用该方案")
else:
    print("\n注意:两种方法结果不一致，建议综合考虑以下因素:")
    print("1. 方法一更注重直接消费差异和运输距离")
    print("2. 方法二考虑了更多经济因素并通过机器学习建模")
    print("3. 可以取两种结果的折中方案或进一步优化权重")



cv = KFold(n_splits=5, shuffle=True, random_state=42)

# 预存每折结果
rmse_ent, rmse_clo = [], []
acc_ent,  acc_clo  = [], []
f1_ent,   f1_clo   = [], []

for train_idx, test_idx in cv.split(scaled_features):
    X_tr, X_te = scaled_features[train_idx], scaled_features[test_idx]
    y_ent_tr, y_ent_te = target_ent.iloc[train_idx], target_ent.iloc[test_idx]
    y_clo_tr, y_clo_te = target_clo.iloc[train_idx], target_clo.iloc[test_idx]

    # 训练并预测
    rf_ent = RandomForestRegressor(n_estimators=150, random_state=42).fit(X_tr, y_ent_tr)
    rf_clo = RandomForestRegressor(n_estimators=150, random_state=42).fit(X_tr, y_clo_tr)
    ent_pred = rf_ent.predict(X_te)
    clo_pred = rf_clo.predict(X_te)

    # 回归指标
    rmse_ent.append(np.sqrt(mean_squared_error(y_ent_te, ent_pred)))
    rmse_clo.append(np.sqrt(mean_squared_error(y_clo_te, clo_pred)))

    # 分类指标（按本折中位数切分）
    th_ent, th_clo = np.median(y_ent_te), np.median(y_clo_te)
    acc_ent.append(accuracy_score(y_ent_te >= th_ent, ent_pred >= th_ent))
    acc_clo.append(accuracy_score(y_clo_te >= th_clo, clo_pred >= th_clo))
    f1_ent.append(f1_score(y_ent_te >= th_ent, ent_pred >= th_ent))
    f1_clo.append(f1_score(y_clo_te >= th_clo, clo_pred >= th_clo))

# 输出 5 折平均结果
print(f"娱乐设施 5 折平均 RMSE: {np.mean(rmse_ent):.2f}")
print(f"服装消费 5 折平均 RMSE: {np.mean(rmse_clo):.2f}")
print(f"娱乐设施 5 折平均 Accuracy: {np.mean(acc_ent):.2f}, F1: {np.mean(f1_ent):.2f}")
print(f"服装消费 5 折平均 Accuracy: {np.mean(acc_clo):.2f}, F1: {np.mean(f1_clo):.2f}")

# 画交叉验证后的柱状图
plt.figure(figsize=(4, 3))
x = np.arange(2)
plt.bar(x - 0.15, [np.mean(acc_ent), np.mean(f1_ent)], 0.3, label='娱乐设施')
plt.bar(x + 0.15, [np.mean(acc_clo), np.mean(f1_clo)], 0.3, label='服装消费')
for i, v in enumerate([np.mean(acc_ent), np.mean(f1_ent)]):
    plt.text(i - 0.15, v + 0.01, f'{v:.2f}', ha='center')
for i, v in enumerate([np.mean(acc_clo), np.mean(f1_clo)]):
    plt.text(i + 0.15, v + 0.01, f'{v:.2f}', ha='center')
plt.xticks(x, ['Accuracy', 'F1-Score'])
plt.ylim(0, 1.1)
plt.title('5 折交叉验证分类性能')
plt.legend()
plt.tight_layout()
plt.show()